"""
邮件模板表结构迁移脚本
"""

from sqlalchemy import create_engine, text
from config import settings

def migrate_email_templates_table():
    """迁移邮件模板表结构"""
    engine = create_engine(settings.database_url)
    
    with engine.connect() as conn:
        # 开始事务
        trans = conn.begin()
        
        try:
            # 检查email_templates表是否存在
            result = conn.execute(text("""
                SELECT COUNT(*) 
                FROM information_schema.tables 
                WHERE table_schema = DATABASE() 
                AND table_name = 'email_templates'
            """))
            
            table_exists = result.scalar() > 0
            
            if table_exists:
                print("邮件模板表已存在，检查是否需要更新字段...")
                
                # 检查是否有新字段
                result = conn.execute(text("""
                    SELECT COUNT(*) 
                    FROM information_schema.columns 
                    WHERE table_schema = DATABASE() 
                    AND table_name = 'email_templates' 
                    AND column_name = 'template_type'
                """))
                
                has_template_type = result.scalar() > 0
                
                if not has_template_type:
                    print("更新邮件模板表字段...")
                    
                    # 添加新字段
                    conn.execute(text("""
                        ALTER TABLE email_templates 
                        ADD COLUMN template_type VARCHAR(20) NOT NULL DEFAULT 'html' COMMENT '模板类型',
                        ADD COLUMN category VARCHAR(50) NOT NULL DEFAULT 'custom' COMMENT '模板分类',
                        ADD COLUMN description TEXT COMMENT '模板描述',
                        ADD COLUMN variables JSON COMMENT '模板变量列表',
                        ADD COLUMN is_active BOOLEAN NOT NULL DEFAULT TRUE COMMENT '是否启用',
                        MODIFY COLUMN name VARCHAR(200) NOT NULL COMMENT '模板名称',
                        MODIFY COLUMN subject VARCHAR(500) NOT NULL COMMENT '邮件主题'
                    """))
                    
                    # 删除旧字段
                    try:
                        conn.execute(text("ALTER TABLE email_templates DROP COLUMN content_type"))
                    except:
                        pass  # 字段可能不存在
                    
                    # 添加索引
                    conn.execute(text("""
                        CREATE INDEX idx_email_templates_user_name ON email_templates(user_id, name)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_email_templates_user_category ON email_templates(user_id, category)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_email_templates_user_active ON email_templates(user_id, is_active)
                    """))
                    conn.execute(text("""
                        CREATE INDEX idx_email_templates_created_at ON email_templates(created_at)
                    """))
                    
                    print("邮件模板表字段更新完成")
                else:
                    print("邮件模板表字段已是最新")
            else:
                print("邮件模板表不存在，将由SQLAlchemy自动创建")
            
            # 提交事务
            trans.commit()
            print("邮件模板表结构迁移完成！")
            
        except Exception as e:
            # 回滚事务
            trans.rollback()
            print(f"迁移失败: {e}")
            raise

if __name__ == "__main__":
    migrate_email_templates_table()
